{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('../')\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from pathlib import Path\n",
    "from BaseTester import BaseTester\n",
    "from demo import u0, f0\n",
    "from MyPlot import *\n",
    "from utils import *\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def kappa_func(xx, yy, mu):\n",
    "\treturn np.exp(-mu * (xx**2 + yy**2))\n",
    "GridSize = 256\n",
    "area=((-1, -1), (1, 1))\n",
    "(left, bottom), (right, top) = area\n",
    "h = (right - left) / (GridSize - 1)\n",
    "\n",
    "xx, yy = np.meshgrid(\n",
    "\tnp.linspace(left, right, GridSize), np.linspace(bottom, top, GridSize)\n",
    ")\n",
    "\n",
    "mus = np.linspace(0, 1, 100)\n",
    "force = np.ones((GridSize, GridSize))\n",
    "# force = normal(xx, yy, h, [(right-left)/2, (top - bottom)/2])\n",
    "\n",
    "sols = []\n",
    "kappas = []\n",
    "for mu in mus:\n",
    "\tkappa = kappa_func(xx, yy, mu) \n",
    "\tA = reaction_A(GridSize, kappa).tocsr()\n",
    "\tb = reaction_b_dir(force, 0, h)\n",
    "\tans = spsolve(A, b).reshape(GridSize, GridSize)\n",
    "\tsols.append(ans)\n",
    "\tkappas.append(kappa)\n",
    "\n",
    "sols = np.stack(sols, axis=0)\n",
    "kappas = np.stack(kappas, axis=0)\n",
    "\n",
    "np.save('DLData/256/Demo2/ValSol.npy', sols)\n",
    "np.save('DLData/256/Demo2/ValKappa.npy', kappas)\n",
    "np.save('DLData/256/Demo2/ValMu.npy', mus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = Path('DLData/256/Demo2/')\n",
    "if not p.is_dir():\n",
    "\tp.mkdir(parents=True)\n",
    "np.save('DLData/256/Demo2/ValSol.npy', sols)\n",
    "np.save('DLData/256/Demo2/ValKappa.npy', kappas)\n",
    "np.save('DLData/256/Demo2/ValMu.npy', mus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResidualLoss(torch.nn.Module):\n",
    "\tdef __init__(self, dtype, device, h, loss_fn, gd=0):\n",
    "\t\tsuper().__init__()\n",
    "\t\tself.dtype = dtype\n",
    "\t\tself.device = device\n",
    "\t\tself.h = h\n",
    "\t\tself.loss_fn = loss_fn\n",
    "\t\tself.gd = gd\n",
    "\n",
    "\t\tself.k = self._get_kernel([[0, -1, 0], [-1, 4, -1], [0, -1, 0]])\n",
    "\t\tself.hard_encode = lambda x: hard_encode(x, self.gd)\n",
    "\n",
    "\tdef _get_kernel(self, k):\n",
    "\t\tk = torch.tensor(k, requires_grad=False)\n",
    "\t\tk = k.view(1, 1, 3, 3).repeat(1, 1, 1, 1).to(self.dtype).to(self.device)\n",
    "\t\treturn k\n",
    "\t\t\n",
    "\tdef forward(self, pre, f, ans):\n",
    "\t\tu = self.hard_encode(pre)\n",
    "\t\tx = F.conv2d(u, self.k)\n",
    "\t\ty = F.conv2d(ans[None, None, ...], self.k) \n",
    "\t\treturn torch.abs(y.squeeze() - x.squeeze())\n",
    "\t\n",
    "class Demo1Tester(BaseTester):\n",
    "\tdef __init__(self, k, **kwargs):\n",
    "\t\tsuper().__init__(**kwargs)\n",
    "\t\tself.img_save_path = kwargs[\"img_save_path\"]\n",
    "\t\tself.k = k\n",
    "\t\tself.mesh()\n",
    "\t\tself.force = f0(self.xx, self.yy, k)\t\t\t\t\t\n",
    "\t\tself.ans = u0(self.xx, self.yy, k)\n",
    "\n",
    "\t\t# self.data = torch.from_numpy(np.stack([self.xx, self.yy])).to(self.dtype).to(self.device)[None, ...]\n",
    "\t\tself.data = torch.ones(1, 1, self.GridSize, self.GridSize).to(self.dtype).to(self.device)\n",
    "\t\t\n",
    "\tdef l2(self, pre, ans):\n",
    "\t\tdiff = (pre - ans) ** 2 * self.h**2\n",
    "\t\tl2_errors = np.sqrt(np.sum(diff))\n",
    "\t\treturn l2_errors.mean() \n",
    "\t\n",
    "\tdef mesh(self):\n",
    "\t\t(self.left, self.bottom), (self.right, self.top) = self.area\n",
    "\t\tself.dx = (self.right - self.left) / (self.GridSize - 1)\n",
    "\t\tself.dy = (self.top - self.bottom) / (self.GridSize - 1)\n",
    "\t\tself.h = self.dx\n",
    "\n",
    "\t\tself.xx, self.yy = np.meshgrid(\n",
    "\t\t\tnp.linspace(self.left, self.right, self.GridSize),\n",
    "\t\t\tnp.linspace(self.bottom, self.top, self.GridSize),\n",
    "\t\t)\n",
    "\t\n",
    "\tdef test(self, exp_names, best_or_last):\n",
    "\t\t# df = {\n",
    "\t\t# \t'exp_name': [],\n",
    "\t\t# \t'l2': []\n",
    "\t\t# }\n",
    "\n",
    "\t\tfor exp_name in exp_names:\n",
    "\t\t\tself.load_kwargs(exp_name)\n",
    "\t\t\tself.load_ckpt(best_or_last, exp_name)\n",
    "\n",
    "\t\t\tls = exp_name.split('-')\n",
    "\t\t\tif 'MatRes' in exp_name or 'EnergyRes' in exp_name:\n",
    "\t\t\t\tact, method = ls[5], ls[-2]\n",
    "\t\t\t\tname = f'res_{act}'\n",
    "\t\t\telse:\n",
    "\t\t\t\tact, method = ls[5], ls[-3]\n",
    "\t\t\t\tname = f'{method.lower()}_{act}'\n",
    "\n",
    "\t\t\twith torch.no_grad():\n",
    "\t\t\t\tpre = self.net(self.data)\n",
    "\t\t\t\tpre = hard_encode(pre, 0).squeeze()\n",
    "\t\t\t\tpre = pre.cpu().numpy()\n",
    "\n",
    "\t\t\t\tans = self.ans\n",
    "\t\t\t\timg_save_path = Path(f\"{self.img_save_path}/{exp_name}\")\n",
    "\t\t\t\tif not img_save_path.is_dir():\n",
    "\t\t\t\t\timg_save_path.mkdir(parents=True)\n",
    "\t\t\t\tself.save_img(f\"{img_save_path}\", pre, ans, name)\n",
    "\n",
    "\t\t\t\tl2_error = self.l2(pre.squeeze(), ans.squeeze())\n",
    "\t\t# \t\tdf['exp_name'].append(exp_name)\n",
    "\t\t# \t\tdf['l2'].append(l2_error)\n",
    "\t\treturn l2_error\n",
    "\t\t# df = pd.DataFrame(df)\n",
    "\t\t# df.to_csv(f\"{self.img_save_path}/k={self.k}-l2.csv\", index=False)\n",
    "\t\n",
    "\tdef save_img(self, save_path, pre, ans, name):\n",
    "\t\tsave_ctf(save_path, pre, ans, self.xx, self.yy, name = name)\n",
    "\t\tlevels = np.linspace(self.ans.min(), self.ans.max(), 12)[2:-2]\n",
    "\t\tsave_contour(save_path, pre, ans, self.xx, self.yy, levels=levels)\n",
    "\n",
    "\t\tbar_ticks = np.linspace(ans.min(), ans.max(), 10)\n",
    "\t\tsave_surf(save_path, pre, self.xx, self.yy, f'surf_{name}')\n",
    "\t\tsave_surf(save_path, ans, self.xx, self.yy, 'surf_ans')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import product\n",
    "\n",
    "df = {\n",
    "\t'K': [],\n",
    "\t'GridSize': [],\n",
    "\t'L2': [],\n",
    "\t'order':[]\n",
    "}\n",
    "print(f\"K\\t&\\tGridSize\\t\\t&\\tl2_error\\t&\\torder\")\n",
    "for k in [1, 2, 4, 8]:\n",
    "# for k in [1]:\n",
    "\tprev_error = None\n",
    "\tprint(prev_error)\n",
    "\tfor GridSize in [32, 64, 128, 256, 512]:\n",
    "\t\tTester = Demo1Tester(\n",
    "\t\t\tk = k,\n",
    "\t\t\tGridSize=GridSize,\n",
    "\t\t\tarea=((-1, -1), (1, 1)),\n",
    "\t\t\tckpt_save_path=f'model_save',\n",
    "\t\t\thyper_parameters_save_path = f'hyper_parameters', \n",
    "\t\t\timg_save_path = f'./images/k={k}', \n",
    "\t\t\tdevice='cuda',\n",
    "\t\t\tdtype=torch.double,\n",
    "\t\t)\n",
    "\t\texp_names = [f'K{k}-{GridSize}-UNet-{GridSize}-Jac-3-100']\n",
    "\t\tl2_error = Tester.test(\n",
    "\t\t\texp_names,\n",
    "\t\t\t'last',\n",
    "\t\t)\n",
    "\t\tif not prev_error is None:\n",
    "\t\t\torder = np.log(prev_error/l2_error) / np.log(2)\n",
    "\t\t\tprev_error = l2_error\n",
    "\t\telse:\n",
    "\t\t\torder = 0\n",
    "\t\t\tprev_error = l2_error\n",
    "\n",
    "\t\tprint(f\"{k}\\t&\\t{GridSize}\\t\\t\\t&\\t{l2_error:.5e}\\t&\\t{order:.5e}\")\n",
    "\t\tdf['K'].append(k)\n",
    "\t\tdf['GridSize'].append(GridSize)\n",
    "\t\tdf['L2'].append(l2_error)\n",
    "\t\tdf['order'].append(order)\n",
    "\t\t\n",
    "df = pd.DataFrame(df)\n",
    "df.to_csv(f\"images/l2.csv\", index=False)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('../')\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from pathlib import Path\n",
    "from BaseTester import BaseTester\n",
    "from MyPlot import *\n",
    "from utils import hard_encode\n",
    "from utils import normal, reaction_A, reaction_b_dir\n",
    "from scipy.sparse.linalg import spsolve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "class Demo2Tester(BaseTester):\n",
    "\tdef __init__(self, testN, **kwargs):\n",
    "\t\tsuper().__init__(**kwargs)\n",
    "\t\tself.testN = testN\n",
    "\t\tself.mesh()\n",
    "\t\t# self.force =  normal(self.xx, self.yy, self.h, ((self.right - self.left)/2, (self.top-self.bottom) / 2))\n",
    "\t\tself.force = np.ones((self.GridSize, self.GridSize))\n",
    "\t\t\n",
    "\t\t\n",
    "\tdef l2(self, pre, ans):\n",
    "\t\tdiff = (pre - ans) ** 2 * self.h**2\n",
    "\t\tl2_errors = np.sqrt(np.sum(diff))\n",
    "\t\treturn l2_errors.mean() \n",
    "\t\n",
    "\tdef mesh(self):\n",
    "\t\t(self.left, self.bottom), (self.right, self.top) = self.area\n",
    "\t\tself.dx = (self.right - self.left) / (self.GridSize - 1)\n",
    "\t\tself.dy = (self.top - self.bottom) / (self.GridSize - 1)\n",
    "\t\tself.h = self.dx\n",
    "\n",
    "\t\tself.xx, self.yy = np.meshgrid(\n",
    "\t\t\tnp.linspace(self.left, self.right, self.GridSize),\n",
    "\t\t\tnp.linspace(self.bottom, self.top, self.GridSize),\n",
    "\t\t)\n",
    "\t\n",
    "\tdef kappa(self, mu):\n",
    "\t\treturn np.exp(-mu * (self.xx ** 2 + self.yy**2))\n",
    "\t\n",
    "\tdef test(self, exp_names, best_or_last, kappas=None):\n",
    "\t\tif kappas is None:\n",
    "\t\t\tmus = np.random.uniform(0, 1,  self.testN)\n",
    "\t\t\tprint(mus)\n",
    "\t\t\tkappas = np.stack([self.kappa(mu)[np.newaxis, ...] for mu in mus ])\n",
    "\t\tans = []\n",
    "\t\tb = reaction_b_dir(self.force, 0, self.h)\n",
    "\t\ttestN = len(kappas)\n",
    "\t\tfor i in tqdm(range(testN)):\n",
    "\t\t\tk = kappas[i].squeeze()\n",
    "\t\t\tA = reaction_A(self.GridSize, k).tocsr()\n",
    "\t\t\tans.append(spsolve(A, b).reshape( self.GridSize, self.GridSize))\n",
    "\t\tans = np.stack(ans)\n",
    "\n",
    "\t\t# np.save('./Demo2-TestAns.npy', ans)\n",
    "\t\t# np.save('./Demo2-TestData.npy', kappas)\n",
    "\t\t\n",
    "\t\t# kappas = np.load('./Demo2-TestData.npy')\n",
    "\t\t# ans = np.load('./Demo2-TestAns.npy')[:, np.newaxis, ...]\n",
    "\t\t# testN = len(kappas)\n",
    "\n",
    "\t\tkappas = kappas[:self.testN]\n",
    "\t\tans = ans[:self.testN]\n",
    "\n",
    "\t\tdata = torch.from_numpy(kappas).to(self.dtype).to(self.device)\n",
    "\t\tfor exp_name in exp_names:\n",
    "\t\t\tdf = {\n",
    "\t\t\t\t# 'test_case': [],\n",
    "\t\t\t\t'l2': []\n",
    "\t\t\t}\n",
    "\t\t\n",
    "\t\t\tself.load_kwargs(exp_name)\n",
    "\t\t\tself.load_ckpt(best_or_last, exp_name)\n",
    "\t\t\n",
    "\t\t\twith torch.no_grad():\n",
    "\t\t\t\tpre = self.net(data)\n",
    "\t\t\t\tpre = hard_encode(pre, 0)\n",
    "\t\t\t\tpre = pre.cpu().numpy().squeeze()\n",
    "\n",
    "\t\t\tprint(pre.shape)\n",
    "\t\t\tsave_path = Path(f\"{self.img_save_path}/{exp_name}\")\n",
    "\t\t\tif not save_path.is_dir():\tsave_path.mkdir(parents=True)\n",
    "\n",
    "\t\t\tl2_errors = []\n",
    "\t\t\tfor i in range(self.testN):\n",
    "\t\t\t\tp = save_path/f\"Case-{i}\"\n",
    "\t\t\t\tif not p.is_dir():\t\n",
    "\t\t\t\t\tp.mkdir(parents=True)\n",
    "\t\t\t\tself.save_img(p, pre[i].squeeze(), ans[i].squeeze(), kappas[i].squeeze())\n",
    "\t\t\t\tl2_error = self.l2(pre[i], ans[i])\n",
    "\t\t\t\tl2_errors.append(l2_error)\n",
    "\t\t\t\t# df['test_case'].append(i)\n",
    "\t\t\t\tdf['l2'].append(l2_error)\n",
    "\t\t\tdf = pd.DataFrame(df)\n",
    "\t\t\tdf.to_csv(f\"{save_path}.csv\")\n",
    "\n",
    "\t\t\tmethod, k, trainN = exp_name.split('-')[-3:]\n",
    "\t\t\tprint(f\"{method}\\t&\\t{k}\\t&\\t{trainN}\\t&\\t{np.array(l2_errors).mean()}\")\n",
    "\t\n",
    "\tdef save_img(self, save_path, pre, ans, cof):\n",
    "\t\tsave_img_force(save_path, cof, name='cof', vmin=0.1, vmax=2)\n",
    "\t\tsave_ctf(save_path, pre, ans, self.xx, self.yy)\n",
    "\t\tlevels = np.linspace(ans.min(), ans.max(), 10)[2:-1]\n",
    "\t\tsave_contour(save_path, pre, ans, self.xx, self.yy, levels=levels)\n",
    "\t\tsave_surf(save_path, pre, self.xx, self.yy, 'surf_pre')\n",
    "\t\tsave_surf(save_path, ans, self.xx, self.yy, 'surf_ans')\n",
    "\t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.04976653 0.91554595 0.7122218 ]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3/3 [00:02<00:00,  1.18it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(3, 256, 256)\n",
      "Jac\t&\t25\t&\t10000\t&\t6.174640490435458e-05\n"
     ]
    }
   ],
   "source": [
    "GridSize = 256\n",
    "Tester = Demo2Tester(\n",
    "\ttestN=3,\n",
    "\tGridSize=GridSize,\n",
    "\tarea=((-1, -1), (1, 1)),\n",
    "\tckpt_save_path=f'Demo2-model_save',\n",
    "\thyper_parameters_save_path = f'Demo2-hyper_parameters', \n",
    "\timg_save_path = f'./images/Demo2-new', \n",
    "\tdevice='cuda',\n",
    "\tdtype=torch.float,\n",
    ")\n",
    "exp_names = ['Demo2-256-Jac-25-10000']\n",
    "# for p in Path(f'model_save/{GridSize}').iterdir():\n",
    "# \tif 'Demo2' in p.stem and '25-5000' in p.stem:\n",
    "# \t\texp_names.append(p.stem)\n",
    "\t\t# print(p.stem.split('-')[-3:])\n",
    "Tester.test(\n",
    "\texp_names,\n",
    "\t'last',\n",
    "\t# kappas=np.ones((0.1, 1, 256, 256)) * 0.2\n",
    "\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Jac\t\t&\t1\t&\t100\t\t&\t0.10053069604839905\n",
    "Jac\t\t&\t1\t&\t500\t\t&\t0.009142189845182969\n",
    "Jac\t\t&\t1\t&\t5000\t&\t0.00037016732885555094\n",
    "\n",
    "Jac\t\t&\t5\t&\t100\t\t&\t0.09415204063292569\n",
    "Jac\t\t&\t5\t&\t500\t\t&\t0.0011334540588817238\n",
    "Jac\t\t&\t5\t&\t5000\t&\t0.00022657807895347414\n",
    "\n",
    "Jac\t\t&\t25\t&\t100\t\t&\t0.017286803618664685\n",
    "Jac\t\t&\t25\t&\t500\t\t&\t0.0006991993553601014\n",
    "Jac\t\t&\t25\t&\t5000\t&\t0.00022067673407101936\n",
    "\n",
    "Desc\t&\t1\t&\t100\t`\t&\t0.09763871757550524\n",
    "Desc\t&\t1\t&\t500\t\t&\t0.0028756982893573682\n",
    "Desc\t&\t1\t&\t5000\t&\t0.0005872470832737366\n",
    "\n",
    "Desc\t&\t5\t&\t100\t\t&\t0.0528270640213064\n",
    "Desc\t&\t5\t&\t500\t\t&\t0.005454670240321949\n",
    "Desc\t&\t5\t&\t5000\t&\t0.000374092234718699\n",
    "\n",
    "Desc\t&\t25\t&\t100\t`\t&\t0.006820678593188237\n",
    "Desc\t&\t25\t&\t500\t\t&\t0.0005588491615975406\n",
    "Desc\t&\t25\t&\t5000\t&\t0.00010590824877410589\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "\n",
    "p = Path('./images/Demo2/256')\n",
    "\n",
    "method = 'Desc'\n",
    "for csv in p.glob(f'*{method}*.csv'):\n",
    "\tK, M = csv.stem.split('-')[-2:]\n",
    "\tdf = pd.read_csv(csv)\n",
    "\tmean = df['l2'].values.mean()\n",
    "\tprint(f\"{method}\\t&\\t{K}\\t&\\t{M}\\t&\\t{mean}\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
